﻿#include "pch.h"
#include <Windows.h>
#include <vector>
#include <iostream>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <string>
#include <chrono>
#include <fstream>
#include <iomanip>
#include <Wincrypt.h>
#include <sstream>

constexpr size_t buffer_size = 4096;

std::string WcharToACP(const wchar_t* wstr)
{
    int len = WideCharToMultiByte(CP_ACP, 0, wstr, -1, NULL, 0, NULL, NULL);
    std::vector<char> buf(len);
    WideCharToMultiByte(CP_ACP, 0, wstr, -1, &buf[0], len, NULL, NULL);
    return std::string(&buf[0]);
}

std::string WcharToUTF8(const wchar_t* wstr)
{
    int len = WideCharToMultiByte(CP_UTF8, 0, wstr, -1, NULL, 0, NULL, NULL);
    std::vector<char> buf(len);
    WideCharToMultiByte(CP_UTF8, 0, wstr, -1, &buf[0], len, NULL, NULL);
    return std::string(&buf[0]);
}

std::string get_current_ip() {
    WSADATA wsa_data;
    if (WSAStartup(MAKEWORD(2, 2), &wsa_data) != 0) {
        std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
        return "";
    }

    char hostname[256];
    if (gethostname(hostname, sizeof(hostname)) != 0) {
        std::cerr << "gethostname failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return "";
    }

    struct addrinfo hints, * res;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC; // 支持 IPv4 和 IPv6
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    if (getaddrinfo(hostname, nullptr, &hints, &res) != 0) {
        std::cerr << "getaddrinfo failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return "";
    }

    sockaddr_in* ipv4;
    sockaddr_in6* ipv6;
    char ip_buffer[INET6_ADDRSTRLEN];
    for (struct addrinfo* ptr = res; ptr != nullptr; ptr = ptr->ai_next) {
        if (ptr->ai_family == AF_INET) { // IPv4
            ipv4 = reinterpret_cast<sockaddr_in*>(ptr->ai_addr);
            inet_ntop(AF_INET, &(ipv4->sin_addr), ip_buffer, sizeof(ip_buffer));
            break;
        }
        else if (ptr->ai_family == AF_INET6) { // IPv6
            ipv6 = reinterpret_cast<sockaddr_in6*>(ptr->ai_addr);
            inet_ntop(AF_INET6, &(ipv6->sin6_addr), ip_buffer, sizeof(ip_buffer));
            break;
        }
    }

    freeaddrinfo(res);
    WSACleanup();

    return std::string(ip_buffer);
}

uint64_t get_current_utc_timestamp_ms() {
    using namespace std::chrono;
    milliseconds ms = duration_cast<milliseconds>(system_clock::now().time_since_epoch());
    return ms.count();
}

std::string calculate_md5(const std::string& file_path) {
    HCRYPTPROV hProv;
    HCRYPTHASH hHash;
    BYTE md5[16];
    DWORD md5_len = sizeof(md5);

    if (!CryptAcquireContext(&hProv, nullptr, nullptr, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT)) {
        std::cerr << "CryptAcquireContext failed: " << GetLastError() << std::endl;
        return "";
    }

    if (!CryptCreateHash(hProv, CALG_MD5, 0, 0, &hHash)) {
        std::cerr << "CryptCreateHash failed: " << GetLastError() << std::endl;
        CryptReleaseContext(hProv, 0);
        return "";
    }

    std::ifstream file(file_path, std::ios::binary);
    if (!file.is_open()) {
        std::cerr << "Unable to open file: " << file_path << std::endl;
        CryptDestroyHash(hHash);
        CryptReleaseContext(hProv, 0);
        return "";
    }

    BYTE buffer[buffer_size];
    size_t bytes_read = 0;

    while (file.read(reinterpret_cast<char*>(buffer), buffer_size) || (bytes_read = file.gcount()) != 0) {
        if (!CryptHashData(hHash, buffer, bytes_read, 0)) {
            std::cerr << "CryptHashData failed: " << GetLastError() << std::endl;
            file.close();
            CryptDestroyHash(hHash);
            CryptReleaseContext(hProv, 0);
            return "";
        }
    }

    file.close();

    if (!CryptGetHashParam(hHash, HP_HASHVAL, md5, &md5_len, 0)) {
        std::cerr << "CryptGetHashParam failed: " << GetLastError() << std::endl;
        CryptDestroyHash(hHash);
        CryptReleaseContext(hProv, 0);
        return "";
    }

    CryptDestroyHash(hHash);
    CryptReleaseContext(hProv, 0);

    std::ostringstream md5_str;
    md5_str << std::hex << std::setfill('0');
    for (size_t i = 0; i < md5_len; ++i) {
        md5_str << std::setw(2) << static_cast<unsigned int>(md5[i]);
    }

    return md5_str.str();
}

std::vector<DriveInfoW> GetDriveInfosW()
{
    // 获取所有盘符的名称
    wchar_t buffer[1024];
    DWORD size = GetLogicalDriveStringsW(sizeof(buffer), buffer);

    std::vector<DriveInfoW> driveInfos;
    for (wchar_t* p = buffer; *p != 0 && p < buffer + size; p += 4)
    {
        DWORD type = GetDriveTypeW(p);
        if (type == DRIVE_FIXED) // 如果是硬盘
        {
            DriveInfoW info;
            info.letter = std::wstring(p);
            info.type = type;

            wchar_t volumeName[MAX_PATH + 1] = { 0 };
            if (GetVolumeInformationW(info.letter.c_str(), volumeName, sizeof(volumeName),
                NULL, NULL, NULL, NULL, 0))
            {
                info.name = volumeName;
            }

            driveInfos.push_back(info);
        }
    }

    return driveInfos;
}

bool CreateDirectories(const std::wstring& path)
{
    std::wstring dirPath = path.substr(0, path.find_last_of(L"/\\"));
    if (dirPath.empty())
    {
        return false;
    }

    if (CreateDirectoryW(dirPath.c_str(), NULL))
    {
        return true;
    }
    else
    {
        DWORD error = GetLastError();
        if (error == ERROR_ALREADY_EXISTS)
        {
            return true;
        }
        else if (error == ERROR_PATH_NOT_FOUND)
        {
            if (CreateDirectories(dirPath))
            {
                return CreateDirectoryW(dirPath.c_str(), NULL);
            }
            else
            {
                return false;
            }
        }
        else
        {
            return false;
        }
    }
}

bool OpenOrCreateFile(const std::wstring& pathAndFile, std::wofstream& file, std::ios_base::openmode _Mode = std::ios_base::out)
{
    std::wstring fullPath = pathAndFile;
    std::wstring path = fullPath.substr(0, fullPath.find_last_of(L"/\\") + 1);
    std::wstring fileName = fullPath.substr(fullPath.find_last_of(L"/\\") + 1);

    file.open(fullPath.c_str(), _Mode);
    if (!file.is_open())
    {
        if (CreateDirectories(path))
        {
            file.open(fullPath.c_str(), std::ios::out | std::ios::app);
        }
    }

    return file.is_open();
}